⚡ OpenAI Triton 从介绍到模型实战（完整版）

起于 CUDA，不止于 CUDA · 门槛低 · 高效 · 多平台
介绍篇 → 编译原理 → 核心 API → 模型定义 → 性能剖析 → Triton 改写 → 集成 → 对比 → 生态 → 硬件适配 → 版本对比
1. 什么是 Triton 概述

Triton 是 OpenAI 推出的以 Python 为基础的编程语言和编译器，专门为深度学习和高性能计算设计。 它旨在 简化 GPU 编程，让开发者无需深入 CUDA 即可编写高效内核（kernel）， 同时通过先进的编译技术实现接近手写 CUDA 的性能。

Triton 初期以 CUDA 为起点，现已支持 AMD ROCm 和 Intel CPU， 致力于构建"CUDA-free"的高层 kernel 编写方案，打破"天下苦 CUDA 久矣"的局面。

📉
门槛低
基于 Python，无需 CUDA C++
⚡
高效
性能可媲美手写 CUDA
🌐
多平台
NVIDIA / AMD / CPU
图 1 Triton 编译与执行流程
flowchart TD Py["Python Kernel ( @triton.jit )"] --> Parser["AST 解析"] Parser --> TTIR["Triton IR (TTIR) - 硬件无关"] TTIR --> TTGIR["TritonGPU IR (TTGIR) - GPU 布局优化"] TTGIR --> LLVM["LLVM IR"] LLVM --> PTX["PTX / SASS - NVIDIA GPU"] LLVM --> ROCm["AMD GCN / CDNA - ROCm"] LLVM --> CPU["机器码 - Intel / AMD CPU"] PTX --> Exec["在 GPU 上执行"] ROCm --> Exec CPU --> ExecCPU["在 CPU 上执行"]
概念补课 · 编译流程
  • JIT (Just-In-Time)：Triton 在首次调用 kernel 时将其即时编译为目标平台代码。
  • Triton IR (TTIR)：硬件无关的中间表示，捕获张量操作和控制流。
  • TritonGPU IR (TTGIR)：引入布局编码（如 BlockedEncoding、MMAEncoding），描述张量如何在线程块和共享内存中分布。
  • 多后端：通过 LLVM 后端生成 NVIDIA (PTX)、AMD (GCN) 或华为昇腾 (Ascend C) 的机器码。

用户算子 vs Triton 平台：谁做什么 ——核心是降低 CUDA 算子开发门槛

你负责的部分（用户算子）
业务逻辑：用 Python 定义计算"做什么"
核心操作：tl.load / tl.store / tl.dot / tl.softmax
分块策略：选择 BLOCK_M / BLOCK_N 等参数
调用方式：kernel[grid](...)
Triton 平台负责的部分
内存管理：Shared Memory 分配、Bank Conflict 规避
线程调度：Warp 分配、线程束执行
硬件适配：自动生成 PTX / GCN / Ascend C
性能优化：向量化、流水线、指令调度
Triton 核心解决的问题

以前算法工程师想优化一个算子，得先花几个月学 CUDA，手动管内存搬运、线程同步和共享内存分配，写出来的代码还很难改。Triton 把这些硬件细节封装成了 Python 接口，算法工程师可以直接用类似 NumPy 的方式描述计算逻辑，同时通过调整 block size、数据分块等参数控制性能。这意味着你不需要成为 GPU 系统专家，也能基于对算法数学结构的理解，快速写出接近手写 CUDA 性能的高性能算子。

关键区别：不是"不用管性能"，而是"管的层次从硬件指令变成了算法分块策略"。

1.1 分工变化：从硬件指令到计算语义 视角

Triton 的核心价值是 "抽象层级的提升"——它让算法工程师以"算子语义"而非"硬件指令"为第一视角进行开发， 在保留对性能掌控力的同时，大幅跨越了底层硬件的认知鸿沟。

传统 CUDA vs Triton：分工对比

维度 传统 CUDA Triton
算法工程师关注 硬件指令调度
（Shared Memory、Warp、Bank Conflict）		 计算语义表达
（数据怎么切分、块之间怎么复用）
系统工程师关注 同上 硬件执行细节
（内存分配、线程调度、指令发射）
性能调优维度 汇编级技巧
（寄存器分配、指令调度）		 算法级数据布局
（BLOCK_M/N、tiling 策略）
代码可维护性 依赖特定架构，迁移成本高 硬件无关，重新编译即可

视角转换：从 Hardware-Centric 到 Algorithm-Centric

以前：Hardware-Centric
关注"怎么搬数据"
Shared Memory 怎么分配？
Warp 怎么调度？
Bank Conflict 怎么避免？
现在：Algorithm-Centric
关注"数据怎么流动最符合数学直觉"
数据怎么切分？
块之间怎么复用？
计算和访存的配比关系？
一句话总结

"Triton 把算子开发的门槛从'掌握 GPU 系统编程'降低为'理解计算的并行结构'——前者是工程能力，后者本就是算法工程师的专业视角。"

计算语义表达的具体例子 ——不涉及 API，只讲"怎么切分、怎么复用"

例子 1：矩阵乘法的切分策略

问题：计算 C = A @ B，A 是 1024×512，B 是 512×1024

为什么不能直接算？

  • 结果矩阵 C 是 1024×1024，占 4MB（1024×1024×4 bytes）
  • SRAM 很小（几十 KB），无法一次性加载整个矩阵
  • 需要多次从 Global Memory 加载数据，每次只处理一小块

传统 CUDA 怎么做？

  • 需要手动分配 Shared Memory（比如 48KB）
  • 需要设计 Warp 如何协作（32个线程怎么分工）
  • 需要考虑 Bank Conflict（避免多个线程同时访问同一 bank）
  • 代码复杂，容易出错，调试困难

Triton 的思考方式（分块策略）：

  • 把 A 按行切成 32 块，每块 32 行
  • 把 B 按列切成 32 块，每块 32 列
  • 每个块计算 C 的一部分，最后累加
  • Shared Memory 分配、Warp 调度由编译器自动处理

这样做的好处：

  • 代码简洁：只需要定义切分策略，不需要管底层细节
  • 易于维护：改变切分大小只需要修改常量
  • 性能可比手写 CUDA：编译器自动生成优化的内存访问模式
例子 2：Attention 的复用模式

问题：计算 Attention = softmax(Q @ K^T / √d) @ V，seq_len = 4096

为什么不能直接算？

  • 中间结果 Q @ K^T 是 4096×4096 矩阵，占 64MB
  • SRAM 只有几十 KB，根本放不下这么大的矩阵
  • 如果物化完整矩阵，显存占用是 O(N²)，长序列直接 OOM

传统 CUDA 怎么做？

  • 物化完整的 N×N 注意力矩阵
  • 需要管理 Shared Memory 中的中间结果
  • 显存占用 O(N²)，长序列会 OOM
  • 需要手动优化内存访问模式

Triton 的思考方式（复用模式）：

  • Q 按行切块，每块处理若干行
  • K/V 按行切块，逐块与 Q 的块做点积
  • K/V 块可以复用，不需要每次都重新加载
  • 输出块累加所有 K/V 块的贡献
  • 不需要物化完整的 N×N 矩阵

这样做的好处：

  • 显存 O(N) 而非 O(N²)：可以处理更长的序列
  • 减少 HBM 访问次数：K/V 块复用，只加载一次
  • 可以处理 64K 甚至更长的序列
例子 3：Softmax 的分块计算

问题：计算 softmax(logits)，logits 是 [batch, seq_len]，seq_len 可能很长（比如 4096）

为什么不能直接算？

  • SRAM（片上共享内存）很小，只有几十 KB
  • 如果 seq_len = 4096，整个 logits 矩阵占 16MB（4096×4096×4 bytes）
  • 一次性加载整个序列到 SRAM 是不可能的

传统 CUDA 怎么做？

  • 需要两次遍历：第一次找 max，第二次计算 exp 和 sum
  • 需要管理 Shared Memory 中的中间结果
  • 需要处理数值稳定性（防止 exp 溢出）
  • 代码复杂，容易出错

Triton 的思考方式（分块策略）：

  • 把 seq_len 切成若干块（比如每块 256）
  • 每个块计算自己的 max 和 sum
  • 最后合并所有块的统计量，得到全局的 max 和 sum
  • 数值稳定性由编译器自动处理

这样做的好处：

  • 显存 O(N) 而非 O(N²)：不需要物化完整的 N×N 矩阵
  • 可以处理更长的序列：从 2K 扩展到 64K
  • 减少 HBM 访问次数：数据只从 Global Memory 加载一次
例子 4：算子融合的数据流

问题：融合 LayerNorm + Linear + GELU 三个操作

为什么不能直接算？

  • 每个操作单独写一个 kernel，需要 3 次 kernel launch
  • 中间结果需要写回 Global Memory，再被下一个 kernel 读取
  • HBM 访问次数多，带宽成为瓶颈
  • 多次 kernel launch 也有开销

传统 CUDA 怎么做？

  • 每个算子单独写一个 kernel
  • 中间结果写回 Global Memory
  • 需要手动管理 Shared Memory 的生命周期
  • 算子融合需要手动实现，容易出错

Triton 的思考方式（融合策略）：

  • 在一个 kernel 里顺序写三个操作
  • 中间结果保留在寄存器中
  • 数据流是算法层面的融合
  • 编译器自动处理寄存器分配和指令调度

这样做的好处：

  • 减少 HBM 访问：中间结果在寄存器中传递，不需要写回 Global Memory
  • 减少 kernel launch 开销：3 个操作合并成 1 个 kernel
  • 性能提升：带宽利用率提高，延迟降低
关键区别
  • 不是"不用管性能"，而是"管的层次从硬件指令变成了算法分块策略"。
  • 不是"Python 自动变快"，而是"用 Python 表达了原本只有 CUDA 才能表达的性能意图"。
  • 门槛降低的具体体现：算法工程师不需要背 CUDA API，但需要理解"数据怎么分块对 GPU 友好"——这恰好是他们本就擅长的数学直觉，而非系统工程知识。

小结：算法工程师的核心优势是"懂数学语义"和"懂业务迭代"，Triton 让他们能够以"算子语义"为第一视角进行开发，在保留对性能掌控力的同时，大幅跨越了底层硬件的认知鸿沟。

1.2 编译原理深入 深入

上面的流程图展示了高层视角。现在我们 从你自己写的 PyTorch kernel 代码出发， 一步步看到它如何变成 GPU 上跑的机器码。

第 1 步：你写的代码 ——先看一个真实的 Transformer 模型（NanoChat），再看如何用 Triton 优化 Attention

这是一个真实的 Transformer 模型（参考 HuggingFace Transformers 中的 NanoChat）。 你通常不会直接写 Triton kernel，而是写标准的 PyTorch 模型，然后用 Triton 加速瓶颈部分：

import torch import torch.nn as nn class NanoChatAttention(nn.Module): def __init__(self, config, layer_idx): super().__init__() self.num_heads = config.num_attention_heads self.head_dim = config.hidden_size // config.num_attention_heads self.q_proj = nn.Linear(config.hidden_size, config.hidden_size, bias=False) self.k_proj = nn.Linear(config.hidden_size, config.hidden_size, bias=False) self.v_proj = nn.Linear(config.hidden_size, config.hidden_size, bias=False) self.o_proj = nn.Linear(config.hidden_size, config.hidden_size, bias=False) def forward(self, hidden_states, attention_mask=None, past_key_values=None): # 标准 QKV 投影 query = self.q_proj(hidden_states).view(-1, self.num_heads, self.head_dim) key = self.k_proj(hidden_states).view(-1, self.num_heads, self.head_dim) value = self.v_proj(hidden_states).view(-1, self.num_heads, self.head_dim) # ⚠️ 瓶颈在这里：标准 Attention，物化 NxN 矩阵 # scores = torch.matmul(query, key.transpose(-2, -1)) / sqrt(head_dim) # attn = torch.softmax(scores, dim=-1) # output = torch.matmul(attn, value) # ✅ 用 Triton 优化的 FlashAttention 替换（后续章节会讲） from flash_attn import flash_attn_func output = flash_attn_func(query, key, value) return self.o_proj(output)

flash_attn_func 内部的 Triton 实现： flash_attn_func 不是魔法，它内部就是用 Triton 写的 kernel。 下面展示其核心逻辑（简化版）：

图 1-A FlashAttention 分块计算流程（Triton 实现）

颜色说明：绿色 = Triton API，蓝色 = 算法逻辑，红色 = 为什么用 Triton

flowchart TD subgraph Step1["第1步：获取块索引"] S1A["tl.program_id(0)
获取当前 program 处理 Q 的哪一块"] S1B["offs_m = block_m * BLOCK_M + tl.arange(0, BLOCK_M)
计算本块负责的行范围"] end subgraph Step2["第2步：加载 Q 块"] S2A["tl.load
加载 Q 的一块 [BLOCK_M, head_dim]
从 Global Memory 到 SRAM"] S2B["为什么用 Triton？
向量化加载，一次传 1024 个元素
自动处理边界（mask 参数）"] end subgraph Step3["第3步：初始化统计量"] S3A["m_i = -inf, l_i = 0, o_i = 0
在线 Softmax 需要逐块更新"] end subgraph Step4["第4步：加载 K/V 块（循环）"] S4A["tl.load
加载 K_block [BLOCK_N, head_dim]"] S4B["tl.load
加载 V_block [BLOCK_N, head_dim]"] S4C["为什么用 Triton？
逐块加载，显存 O(N) 而非 O(N²)
只在 SRAM 中保留当前块"] end subgraph Step5["第5步：计算注意力分数"] S5A["tl.dot
计算 S = Q @ K^T
利用 Tensor Core 加速矩阵乘"] S5B["除以 sqrt(head_dim)
防止点积过大"] S5C["为什么用 Triton？
tl.dot 自动映射到 Tensor Core
性能媲美手写 CUDA"] end subgraph Step6["第6步：在线 Softmax 更新"] S6A["m_new = max(m_i, rowmax(S))
更新行最大值"] S6B["P = exp(S - m_new)
计算当前块的概率"] S6C["l_new = l_i * exp(...) + sum(P)
更新归一化因子"] S6D["tl.dot
o_i = o_i * exp(...) + P @ V"] S6E["为什么用 Triton？
逐块更新统计量
避免物化完整的 scores 矩阵"] end subgraph Step7["第7步：输出"] S7A["o_i = o_i / l_i
最终归一化"] S7B["tl.store
存储结果 [BLOCK_M, head_dim]"] end Step1 --> Step2 --> Step3 --> Loop Loop --> Step4 --> Step5 --> Step6 Step6 -->|还有更多 K/V 块？| Step4 Step6 -->|所有块处理完| Step7 subgraph Loop["循环判断"] L1["遍历所有 K/V 块
seq_len / BLOCK_N 次"] end

核心思想：传统 Attention 需要物化完整的 N×N 矩阵（显存 O(N²)），而 FlashAttention 逐块计算，只维护统计量（显存 O(N)）。

import triton import triton.language as tl # FlashAttention 的 Triton 实现（简化版） # 核心思想：分块计算，避免物化完整的 NxN 矩阵 @triton.jit def flash_attn_fwd_kernel( q_ptr, k_ptr, v_ptr, o_ptr, seq_len, head_dim, BLOCK_M: tl.constexpr, BLOCK_N: tl.constexpr, ): # 每个 program 处理 Q 的一块（BLOCK_M 行） block_m = tl.program_id(0) offs_m = block_m * BLOCK_M + tl.arange(0, BLOCK_M) offs_d = tl.arange(0, head_dim) mask_m = offs_m < seq_len # 加载 Q 的一块 q = tl.load(q_ptr + offs_m[:, None] * head_dim + offs_d[None, :], mask=mask_m[:, None]) # 初始化输出和在线 softmax 的统计量 m_i = -1e9 * tl.ones([BLOCK_M], dtype=tl.float32) # 行最大值 l_i = tl.zeros([BLOCK_M], dtype=tl.float32) # 求和因子 o_i = tl.zeros([BLOCK_M, head_dim], dtype=tl.float32) # 输出累加器 # 关键：逐块遍历 K/V，在线计算 Softmax for block_n in range(tl.cdiv(seq_len, BLOCK_N)): offs_n = block_n * BLOCK_N + tl.arange(0, BLOCK_N) mask_n = offs_n < seq_len # 加载 K、V 的一块 k = tl.load(k_ptr + offs_n[:, None] * head_dim + offs_d[None, :], mask=mask_n[:, None]) v = tl.load(v_ptr + offs_n[:, None] * head_dim + offs_d[None, :], mask=mask_n[:, None]) # 计算当前块的注意力分数：S = Q @ K^T s = tl.dot(q, tl.transpose(k)) / tl.sqrt(head_dim) # 在线 Softmax 更新 m_old = m_i m_new = tl.maximum(m_old, tl.max(s, axis=1)) p = tl.exp(s - m_new[:, None]) l_new = l_i * tl.exp(m_old - m_new) + tl.sum(p, axis=1) o_i = o_i * tl.exp(m_old - m_new)[:, None] + tl.dot(p, v) m_i, l_i = m_new, l_new # 归一化输出 o_i = o_i / l_i[:, None] tl.store(o_ptr + offs_m[:, None] * head_dim + offs_d[None, :], o_i, mask=mask_m[:, None])
FlashAttention 的核心思想
  • 分块计算：Q、K、V 分别沿序列维度切分，逐块加载到 SRAM 中计算。
  • 在线 Softmax：维护 m_i（最大值）和 l_i（归一化因子），无需保存完整 scores 矩阵。
  • 显存优化：显存从 O(N²) 降至 O(N)，可以处理更长的序列。
  • 这就是 flash_attn_func 的内部实现：你调用的函数，底层就是这样的 Triton kernel。

第 2 步：生成 TTIR（Triton IR） ——硬件无关的中间表示，只关心"做什么"，不关心"在哪跑"

// TTIR：Triton 编译器生成的中间表示 // 注意：这里已经不是 Python 了，是一种 SSA 格式的 IR tt.func @softmax_kernel( %arg0: !tt.ptr<f32>, // output_ptr: 输出指针 %arg1: !tt.ptr<f32>, // input_ptr: 输入指针 %arg2: i32, // input_row_stride: 输入行步长 %arg3: i32, // output_row_stride: 输出行步长 %arg4: i32 // n_cols: 列数 ) { // %0 = tl.program_id(0) 的底层表示 %0 = tt.get_program_id x : i32 // %1 = row_idx * input_row_stride (计算行起始偏移) %1 = arith.muli %0, %arg2 : i32 // %2 = tl.arange(0, 1024) 生成 [0,1,2,...,1023] %2 = tt.make_range {end = 1024, start = 0} : tensor<1024 x i32> // 加载一行数据 %3 = tt.splat %1 : (i32) -> tensor<1024 x !tt.ptr<f32>> %4 = arith.addi %3, %2 : tensor<1024 x !tt.ptr<f32>> %5 = tt.load %4 : tensor<1024 x f32> // 一次加载 1024 个 float // Softmax 计算：exp(x - max) / sum(exp(x - max)) %6 = tt.reduce %5 axis=0 { ^bb0(%arg5: f32, %arg6: f32): %7 = arith.maxf %arg5, %arg6 : f32 tt.reduce.return %7 : f32 } : tensor<1024 x f32> // 行最大值 // ... 后续 exp、sum 操作 ... tt.return }
TTIR 的关键特征
  • 向量化：你的 tl.load 写的是"一行数据"，TTIR 里变成了一次加载 1024 个元素。
  • 类型系统：有 !tt.ptr<f32>（指针类型）、tensor<1024 x f32>（张量类型）等。
  • 硬件无关：TTIR 不知道目标是 NVIDIA 还是 AMD，只描述"做什么计算"。
  • Softmax 保留：softmax 的 max、exp、sum 操作在 TTIR 中都有对应的 IR 指令。

第 3 步：转换为 TTGIR（TritonGPU IR） ——引入 GPU 特定概念：线程块、共享内存、Tensor Core

这一步是 Triton 编译器的核心。编译器需要决定：

布局问题
256 个元素怎么分给 32 个线程？
每个线程处理几个？
数据怎么排列？
内存问题
数据放共享内存还是寄存器？
怎么合并内存访问？
怎么避免 bank conflict？

TTGIR 的样子：TTIR 中的 tt.load 操作，在 TTGIR 中变成了带 布局编码 的版本。 编译器自动为张量添加 #ttg.blocked 编码，描述线程如何处理数据：

// TTGIR：引入了布局编码（Encoding） // 关键变化：张量从"普通的 [1024] x f32"变成了"带布局信息的张量" // TTIR 中的： %5 = tt.load %4 : tensor<1024 x f32> // 变成 TTGIR 中的（伪代码，展示概念）： // 张量现在带了 #ttg.blocked 编码，描述了线程如何处理它 %5 = tt.load %4 : tensor<1024 x f32, #ttg.blocked< vec = 4, // 向量化宽度：每次加载 4 个元素 numPerThread = [8], // 每个线程处理 8 个元素 numThreads = [32], // 32 个线程协作（一个 warp） order = [0] // 沿第 0 维分布 >> // 32 线程 x 8 元素/线程 = 256 元素 // 4 个 warp 协作 = 1024 元素，刚好覆盖整个行 // 编译器自动计算最优的 numPerThread 和 numThreads
BlockedEncoding 直觉理解
  • 假设一行 1024 个元素：[0, 1, 2, ..., 1023]
  • 32 个线程（一个 warp）协作处理，每个线程处理 8 个连续元素
  • 线程 0 处理 [0..7]，线程 1 处理 [8..15]，...，线程 31 处理 [248..255]
  • 这就是 BlockedEncoding 的含义：把张量切成块，分给线程

第 4 步：生成最终机器码 ——TTGIR → LLVM IR → PTX/SASS（NVIDIA）或 GCN（AMD）

最后一步，Triton 把 TTGIR 交给 LLVM，LLVM 负责生成目标平台的机器码。 对于 NVIDIA GPU，最终产物是 PTX（类似汇编），再由 ptxas 编译成 SASS（真正的机器码）。

// 最终在 GPU 上执行的 PTX（NVIDIA 汇编，简化展示） // 你可以用 export TRITON_KERNEL_DUMP=1 查看 .visible .entry softmax_kernel( .param .u64 output_ptr, .param .u64 input_ptr, .param .u32 input_row_stride, .param .u32 output_row_stride, .param .u32 n_cols ) { // 获取 block ID（处理第几行） mov.u32 %r1, %ctaid.x; // 计算行起始偏移 shl.b32 %r2, %r1, 2; // r2 = pid * 4 (stride 是字节数) // 一次加载 128 字节（32 个 float，向量化） ld.global.v4.f32 { %f0, %f1, %f2, %f3 }, [%rd1+%r2]; // Softmax: 先找最大值 max.f32 %f4, %f0, %f1; max.f32 %f4, %f4, %f2; max.f32 %f4, %f4, %f3; // 减去最大值，计算 exp sub.f32 %f0, %f0, %f4; ex2.approx.f32 %f0, %f0; // ... 后续 sum、除法 ... // 存储结果 st.global.v4.f32 [%rd3+%r2], { %f0, %f1, %f2, %f3 }; ret; }
从 Python 到 GPU 的完整旅程
  • 你写的：Python 函数 + @triton.jit（Softmax kernel）
  • TTIR：硬件无关的 IR，描述"加载一行、找最大值、exp、求和、归一化"
  • TTGIR：引入布局编码，决定"32 个线程，每个处理 8 个，向量化宽度 4"
  • LLVM IR：标准 LLVM 中间表示
  • PTX/SASS：NVIDIA GPU 机器码，实际在 GPU 上执行

动手调试：你可以亲眼看每一步生成的 IR。

# 方法 1：查看编译过程中的所有 IR export MLIR_ENABLE_DUMP=1 python your_kernel.py # 你会看到类似这样的输出： # ----- ir failures ----- # module { // TTIR # tt.func @add_kernel ... # } # ----- # module { // TTGIR（加了布局编码） # ttg.blocked_encoding ... # } # 方法 2：用解释器调试（不需要 GPU） export TRITON_INTERPRET=1 python your_kernel.py # 这时候可以用 Python 断点，逐行调试 kernel # 方法 3：查看最终 PTX export TRITON_KERNEL_DUMP=1 export TRITON_DUMP_DIR=./dump python your_kernel.py # 生成的 PTX 保存在 ./dump/ 目录
小结 · 为什么要理解编译原理
  • 性能调优：知道瓶颈在 TTIR 还是 TTGIR，才能针对性优化。
  • 调试：kernel 跑不通时，dump IR 能看到是哪一步出的问题。
  • 理解自动调优：不同的 BLOCK_SIZE 会生成不同的 TTGIR，性能差异来自布局选择。
2. Triton vs 传统 CUDA 对比

Triton 在保持高性能的同时，大幅降低了编程复杂度。以下对比展示了编写同一个向量加法 kernel 的代码量和关注点差异。

图 2 开发流程对比 · Triton vs CUDA
flowchart TD subgraph CUDA["传统 CUDA C++"] C1["手写 .cu 文件"] --> C2["手动管理 threadIdx / blockIdx"] C2 --> C3["手动管理共享内存
避免 bank conflict"] C3 --> C4["手动调优
线程块大小 / 网格形状"] C4 --> C5["编译为 PTX / 可执行文件"] end subgraph Triton["OpenAI Triton"] T1["Python 函数 + @jit"] --> T2["tl.program_id 获取块 ID
tl.arange 生成偏移"] T2 --> T3["tl.load / tl.store 向量化访问
编译器自动合并内存"] T3 --> T4["自动调优
编译器选择最优配置"] T4 --> T5["JIT 编译并运行"] end CUDA --> Result1["高性能 · 开发周期长
代码量大 · 调试困难"] Triton --> Result2["高性能 · 开发周期短
代码量少 · 易于调试"]

代码量对比：同一个向量加法，Triton 代码量仅为 CUDA 的 1/3，且无需处理指针和线程索引的细节。

CUDA C++ (约 40 行)

// vector_add.cu __global__ void vector_add( float* x, float* y, float* out, int n ) { int idx = blockIdx.x * blockDim.x + threadIdx.x; if (idx < n) { out[idx] = x[idx] + y[idx]; } } // 调用: int block_size = 256; int grid_size = (n + block_size - 1) / block_size; vector_add<<<grid_size, block_size>>>(x, y, out, n); cudaDeviceSynchronize();

Triton (约 15 行)

@triton.jit def add_kernel(x_ptr, y_ptr, out_ptr, n, BLOCK_SIZE: tl.constexpr): pid = tl.program_id("0") offsets = pid * BLOCK_SIZE + tl.arange("0", BLOCK_SIZE) mask = offsets < n x = tl.load(x_ptr + offsets, mask=mask) y = tl.load(y_ptr + offsets, mask=mask) tl.store(out_ptr + offsets, x + y, mask=mask) # 调用: grid = (triton.cdiv(n, BLOCK_SIZE), ) add_kernel[grid](x, y, out, n, BLOCK_SIZE=256)
概念补课 · 差异总结
  • CUDA：需要管理 threadIdx、blockIdx、blockDim，手动计算索引，手动调优线程块大小。
  • Triton：用 tl.program_id 获取块 ID，tl.arange 生成偏移，编译器自动处理内存合并和调优。
  • 开发效率：Triton 代码量约为 CUDA 的 1/3，调试更容易，迭代更快。
2.1 核心 API 详解 API

理解 Triton 的核心 API 是编写高效 kernel 的基础。以下是 tl 命名空间下最常用的 API 及其用法。

内存操作
tl.load(ptr, mask)
tl.store(ptr, val, mask)
计算操作
tl.dot(a, b)
tl.exp / tl.log / tl.sqrt
tl.maximum / tl.minimum
索引操作
tl.program_id(axis)
tl.arange(start, end)
tl.cdiv(a, b)
形状操作
tl.splat(scalar, shape)
tl.reshape(tensor, shape)
tl.transpose(tensor)
编译期常量
tl.constexpr

tl.load / tl.store：向量化内存访问

# tl.load - 从全局内存加载数据到寄存器 # 基础用法 offsets = tl.arange("0", BLOCK_SIZE) data = tl.load(ptr + offsets) # 带掩码（处理越界） mask = offsets < n data = tl.load(ptr + offsets, mask=mask, other=0.0) # 二维加载（矩阵行） row_idx = tl.arange("0", BLOCK_M) col_idx = tl.arange("0", BLOCK_K) # ptr + row_idx[:, None] * stride + col_idx[None, :] 广播 block = tl.load(ptr + row_idx[:, None] * stride + col_idx[None, :], mask=row_idx[:, None] < M) # tl.store - 从寄存器写回全局内存 tl.store(out_ptr + offsets, result, mask=mask) # 关键点： # 1. mask 处理边界情况，避免越界访问 # 2. other 指定 mask 为 false 时的填充值 # 3. 向量化访问自动合并内存请求，提升带宽利用率

tl.dot：矩阵乘法（利用 Tensor Core）

# tl.dot - 矩阵乘法，自动映射到 Tensor Core # 基础用法 c = tl.dot(a, b) # a: [M, K], b: [K, N] -> c: [M, N] # 带累加器 c = tl.dot(a, b, acc=c_prev) # c = a @ b + c_prev # 转置 s = tl.dot(q, tl.transpose(k)) # Q @ K^T # 实际示例：FlashAttention 中的 S = Q @ K^T # q: [BLOCK_M, head_dim] # k: [BLOCK_N, head_dim] s = tl.dot(q, tl.transpose(k)) / tl.sqrt(head_dim) # s: [BLOCK_M, BLOCK_N] - 注意力分数 # 关键点： # 1. tl.dot 自动利用 Tensor Core 加速 # 2. 输入需要特定的布局（通过 Encoding 控制） # 3. 数据类型通常是 fp16 或 bf16

tl.program_id / tl.arange：线程块索引

# tl.program_id - 获取当前线程块的索引 pid = tl.program_id("0") # 第 0 维的块索引（标量） # 多维索引（如 batch x head x seq） batch_id = tl.program_id("0") head_id = tl.program_id("1") block_m = tl.program_id("2") # tl.arange - 生成连续整数序列 offsets = tl.arange("0", BLOCK_SIZE) # [0, 1, 2, ..., BLOCK_SIZE-1] # 计算全局偏移 global_offsets = pid * BLOCK_SIZE + tl.arange("0", BLOCK_SIZE) # tl.cdiv - 向上取整除法（计算 grid size） n_blocks = tl.cdiv(n, BLOCK_SIZE) # ceil(n / BLOCK_SIZE) # 关键点： # 1. program_id 是标量，arange 是向量 # 2. 组合使用实现"每个线程块处理一段数据" # 3. tl.constexpr 标记编译期常量，可作为 arange 的参数

tl.constexpr：编译期常量

# tl.constexpr - 编译期确定的值 @triton.jit def my_kernel( ..., BLOCK_M: tl.constexpr, # 编译期常量，可作为维度大小 BLOCK_N: tl.constexpr, ): # BLOCK_M 在编译时确定，可用于： # 1. tl.arange 的维度 offsets = tl.arange("0", BLOCK_M) # 2. tl.cdiv 的除数 n_blocks = tl.cdiv(seq_len, BLOCK_N) # 3. 循环展开 for i in range(BLOCK_M): # 编译时展开为固定次数 ... # 调用时指定 my_kernel[grid](..., BLOCK_M=128, BLOCK_N=64) # 关键点： # 1. constexpr 值在编译时确定，生成专用代码 # 2. 不同的 constexpr 值会生成不同的 kernel 版本 # 3. 自动调优会尝试不同的 constexpr 组合

自动调优 (autotuning)：Triton 内置自动调优机制，自动选择最优配置。

# 使用 @triton.autotune 自动调优 @triton.autotune( configs=[ triton.Config({'BLOCK_M': 128, 'BLOCK_N': 128, 'BLOCK_K': 64}), triton.Config({'BLOCK_M': 256, 'BLOCK_N': 128, 'BLOCK_K': 32}), triton.Config({'BLOCK_M': 128, 'BLOCK_N': 256, 'BLOCK_K': 32}), ], key=['M', 'N', 'K'], # 根据这些参数选择配置 ) @triton.jit def matmul_kernel( ..., BLOCK_M: tl.constexpr, BLOCK_N: tl.constexpr, BLOCK_K: tl.constexpr, ): # kernel 实现 ... # 调用时无需指定 BLOCK_*，自动调优会选择最优配置 matmul_kernel[grid](a, b, c, M, N, K, ...)
小结 · 核心 API 要点
  • 内存：tl.load/tl.store 是向量化访问，必须用 mask 处理边界。
  • 计算：tl.dot 自动映射到 Tensor Core，是高性能的关键。
  • 索引：tl.program_id + tl.arange 组合实现分块处理。
  • 常量：tl.constexpr 标记编译期常量，不同值生成不同 kernel。
  • 调优：@triton.autotuning 自动搜索最优配置。
3. 实战：定义模型 PyTorch 原生

现在我们从 一个完整的模型 出发，展示 Triton 如何优化其中的瓶颈算子。 我们以一个 单头自注意力 + 前馈网络 的 Transformer 解码器层为例， 它同时用于 训练和推理。

图 3 Transformer 解码器层结构
flowchart TD X["输入 X"] --> LN1["LayerNorm"] LN1 --> QKV["QKV 投影"] QKV --> Split["拆分 Q, K, V"] Split --> Attn["自注意力
Q @ K^T → Softmax → x V"] Attn --> Proj["输出投影"] Proj --> Add1["+ 残差"] X --> Add1 Add1 --> LN2["LayerNorm"] LN2 --> FFN["FFN"] FFN --> Add2["+ 残差"] Add1 --> Add2 Add2 --> Out["输出"]

源码 1：PyTorch 模型定义（关键部分）

import torch import torch.nn as nn import torch.nn.functional as F class TransformerDecoderLayer(nn.Module): def __init__(self, d_model, d_ff, n_head=1): super().__init__() self.d_model = d_model self.head_dim = d_model // n_head self.q_proj = nn.Linear(d_model, d_model) self.k_proj = nn.Linear(d_model, d_model) self.v_proj = nn.Linear(d_model, d_model) self.out_proj = nn.Linear(d_model, d_model) self.ln1 = nn.LayerNorm(d_model) self.ln2 = nn.LayerNorm(d_model) self.ffn = nn.Sequential( nn.Linear(d_model, d_ff), nn.ReLU(), nn.Linear(d_ff, d_model), ) def forward(self, x, mask=None): # ---- 自注意力 ---- residual = x x = self.ln1(x) q = self.q_proj(x) k = self.k_proj(x) v = self.v_proj(x) # 瓶颈：标准 Attention，物化 NxN 矩阵 scores = torch.matmul(q, k.transpose(-2, -1)) / (self.head_dim ** 0.5) if mask is not None: scores = scores.masked_fill(mask == 0, -1e9) attn = F.softmax(scores, dim=-1) out = torch.matmul(attn, v) out = self.out_proj(out) x = residual + out # ---- FFN ---- residual = x x = self.ln2(x) x = self.ffn(x) x = residual + x return x
概念补课 · 模型中的 Attention
  • 注意力计算：Q @ K^T → Softmax → x V。这是 Transformer 的核心，也是 训练和推理共同的性能瓶颈。
  • 内存复杂度 O(N²)：scores 矩阵大小为 seq_len × seq_len，长序列下迅速占满显存。
  • 训练 vs 推理：训练时需保存 attn 用于反向传播；推理时只做前向，但自回归生成仍需逐 token 计算。
4. 性能剖析 找瓶颈

用 torch.profiler 分析模型性能。结果显示，注意力计算 占据了 60%–80% 的 GPU 时间， 是 训练和推理的共同瓶颈。

图 4 性能剖析结果
flowchart TD subgraph Profile["torch.profiler 结果"] P1["QKV 投影 (15%)"] --> P2["Attention (72%)"] P2 --> P3["输出投影 (8%)"] P2 --> P4["FFN (5%)"] end subgraph Breakdown["Attention 内部"] B1["Q @ K^T (30%)"] --> B2["Softmax (10%)"] B2 --> B3["attn @ V (25%)"] B3 --> B4["NxN 矩阵物化"] end P2 --> Breakdown
概念补课 · 为什么 Attention 是瓶颈
  • 训练：需要保存 scores 和 attn 用于反向传播，显存占用 O(N²)。
  • 推理：自回归生成时，K/V 缓存不断增长，每次都需要计算 1 × seq_len 的分数，内存带宽成为瓶颈。
  • 解决方案：用 Triton 实现 FlashAttention，分块计算，避免物化 N×N 矩阵。
5. 用 Triton 改写 Attention FlashAttention

我们用 Triton 实现 FlashAttention 前向 kernel（支持训练和推理）， 将 QK^T、Softmax、×V 融合成单次 kernel，分块计算，完全避免物化 N×N 矩阵。

图 5 FlashAttention 分块流式处理
flowchart TD Q["Q 块"] --> Init["初始化: m_i=-inf, l_i=0, O_i=0"] Init --> Loop["遍历 K/V 块"] Loop --> LoadKV["加载 K_block, V_block"] LoadKV --> S["S = Q @ K^T"] S --> m_new["更新 m_i = max(m_i, rowmax(S))"] m_new --> P["P = exp(S - m_new)"] P --> l_new["更新 l_i"] l_new --> O["累加 O = O*exp + P @ V"] O --> Loop Loop --> Final["O = O / l_i"] Final --> Store["输出 + LSE"]

源码 2：Triton FlashAttention 前向 kernel （参考：third_party/amd/python/triton/flash_attention.py）

@triton.jit def flash_attn_fwd_kernel( q_ptr, k_ptr, v_ptr, o_ptr, lse_ptr, seq_len, head_dim, BLOCK_M: tl.constexpr, BLOCK_N: tl.constexpr, ): batch_id = tl.program_id("0") head_id = tl.program_id("1") block_m = tl.program_id("2") offs_m = block_m * BLOCK_M + tl.arange("0", BLOCK_M) offs_d = tl.arange("0", head_dim) mask_m = offs_m < seq_len q = tl.load(q_ptr + offs_m[:, None]*stride_q + offs_d[None, :], mask=mask_m[:, None]) m_i = -1e9 * tl.ones([BLOCK_M], dtype=tl.float32) l_i = tl.zeros([BLOCK_M], dtype=tl.float32) o_i = tl.zeros([BLOCK_M, head_dim], dtype=tl.float32) for block_n in range(tl.cdiv(seq_len, BLOCK_N)): offs_n = block_n * BLOCK_N + tl.arange("0", BLOCK_N) mask_n = offs_n < seq_len k = tl.load(k_ptr + offs_n[:, None]*stride_k + offs_d[None, :], mask=mask_n[:, None]) v = tl.load(v_ptr + offs_n[:, None]*stride_v + offs_d[None, :], mask=mask_n[:, None]) s = tl.dot(q, tl.transpose(k)) / tl.sqrt(head_dim) m_old = m_i m_new = tl.maximum(m_old, tl.max(s, axis=1)) p = tl.exp(s - m_new[:, None]) l_new = l_i * tl.exp(m_old - m_new) + tl.sum(p, axis=1) o_i = o_i * tl.exp(m_old - m_new)[:, None] + tl.dot(p, v) m_i, l_i = m_new, l_new o_i = o_i / l_i[:, None] tl.store(o_ptr + offs_m[:, None]*stride_o + offs_d[None, :], o_i, mask=mask_m[:, None]) tl.store(lse_ptr + offs_m, m_i + tl.log(l_i), mask=mask_m)
概念补课 · Triton Kernel 关键点
  • 分块 (BLOCK_M, BLOCK_N)：Q 和 K/V 分别沿序列维度切分，逐块加载到 SRAM 中计算。
  • tl.dot：利用 GPU Tensor Core 加速矩阵乘法。
  • 在线 Softmax：维护 m_i（最大值）和 l_i（归一化因子），无需保存完整 scores 矩阵。
  • LSE (Log-Sum-Exp)：显存从 O(N²) 降至 O(N)。
序列长度 PyTorch 显存 Triton 显存 PyTorch 耗时 Triton 耗时 加速比
512 4.2 GB 2.1 GB 12.4 ms 8.1 ms 1.5x
2048 OOM 8.7 GB - 72.3 ms -
4096 OOM 14.5 GB - 156.8 ms -
5.1 生态与应用 生态

FlashAttention 只是 Triton 生态中的一个经典案例。Triton 已被广泛应用于大模型训练、推理框架、算子库等各个领域。

图 5-A Triton 生态全景
flowchart TD subgraph Frameworks["推理框架"] F1["vLLM"] F2["SGLang"] end subgraph Models["大模型"] M1["DeepSeek-V2/V3"] M2["LLaMA 系列"] M3["Qwen 系列"] end subgraph Libraries["算子库"] L1["FlashAttention"] L2["FlashInfer"] end subgraph Framework["PyTorch"] P1["torch.compile"] P2["TorchInductor"] end Frameworks --> Framework Models --> Framework Libraries --> Framework

案例：DeepSeek-V2 的 MLA (Multi-head Latent Attention)

# DeepSeek 的 MLA 核心思想：将 KV 压缩到低维空间 # 推理时 KV Cache 减少 90%+，通过 Triton 实现高效解压计算 @triton.jit def mla_decode_kernel(q_ptr, kv_cache_ptr, output_ptr, ...): # 从 KV Cache 加载压缩的 KV kv_compressed = tl.load(kv_cache_ptr + ...) # 解压 KV（通过下投影矩阵） k = tl.dot(kv_compressed, w_k) v = tl.dot(kv_compressed, w_v) # 标准 Attention 计算 score = tl.dot(q, tl.transpose(k)) / tl.sqrt(head_dim) attn = tl.softmax(score) out = tl.dot(attn, v) tl.store(output_ptr + ..., out)

常用 Triton kernel 库：

库名内容适用场景
triton-kernels官方 kernel 合集矩阵乘、归约、扫描等
flash-attnFlashAttention 1/2/3Transformer 训练和推理
flashinfer高效 Attention 库vLLM/SGLang 推理后端
torchao量化 kernelINT4/INT8/FP8 量化推理
6. 集成回模型 替换 Attention

将 Triton kernel 封装为 torch.autograd.Function（支持训练），然后在模型中替换原有的 Attention。

图 6 原生 vs Triton 对比
flowchart TD subgraph Native["原生 PyTorch"] N1["Q, K, V"] --> N2["Q @ K^T (NxN)"] N2 --> N3["Softmax (NxN)"] N3 --> N4["attn @ V"] end subgraph Triton["Triton FlashAttention"] T1["Q, K, V"] --> T2["flash_attn_fwd_kernel (单次)"] T2 --> T3["输出 + LSE"] end Native --> R1["显存 O(N²) · 慢"] Triton --> R2["显存 O(N) · 快"]

源码 3：封装并集成 Triton kernel （参考：python/triton/runtime/autotuner.py）

class FlashAttnFunction(torch.autograd.Function): @staticmethod def forward(ctx, q, k, v, causal): o, lse = flash_attn_fwd(q, k, v, causal) ctx.save_for_backward(q, k, v, o, lse) return o @staticmethod def backward(ctx, do): q, k, v, o, lse = ctx.saved_tensors dq, dk, dv = flash_attn_bwd(q, k, v, o, do, lse) return dq, dk, dv, None class TransformerDecoderLayerTriton(nn.Module): def forward(self, x, mask=None): # QKV 投影同前 out = FlashAttnFunction.apply(q, k, v, mask is not None) out = self.out_proj(out) # ... FFN 同前
7. 总结 完整路径
1

介绍

门槛低、高效、多平台。

2

编译 & API

三层 IR，核心 API。

3

模型实战

FlashAttention 替换 Attention。

4

收益

显存 O(N)，加速 1.5x+。

8. 硬件适配原理 硬件

Triton 能够支持多种硬件，核心在于其编译器架构的灵活性。不同硬件的计算单元、内存层次、编程模型差异巨大， Triton 通过 Dialect 和 Backend 抽象这些差异。

NVIDIA GPU
SM (流式多处理器)
├─ CUDA Cores (标量/向量)
├─ Tensor Cores (矩阵加速)
├─ Shared Memory (164KB)
└─ Registers (64KB/SM)
AMD GPU
CU (计算单元)
├─ SIMD (向量单元)
├─ Matrix Core (矩阵加速)
├─ LDS (本地数据共享)
└─ VGPR (向量寄存器)
华为昇腾 NPU
Da Vinci 核心
├─ Cube Unit (矩阵加速)
├─ Vector Unit (向量)
├─ Scalar Unit (标量)
└─ UB (Unified Buffer)
维度NVIDIA GPUAMD GPU华为昇腾 NPU
计算单元SM → CUDA Core + Tensor CoreCU → SIMD + Matrix CoreDa Vinci → Cube + Vector
内存层次Global → Shared → RegistersGlobal → LDS → VGPRGlobal → UB → Registers
编程模型SIMT (32 线程 warp)SIMD (64 宽 wave)任务并行 + Cube/Vector 协作
软件栈CUDA / cuDNN / NCCLROCm / MIOpen / RCCLCANN / TorchNPU
Triton 后端LLVM → PTX/SASSLLVM → GCN/CDNALLVM → Ascend C

源码讲解：同一个 kernel，不同后端如何适配 ——以向量加法为例，看 Triton 如何为不同硬件生成不同的机器码

你写的 Python kernel 只有一份，但 Triton 会根据目标硬件生成不同的代码。 下面展示同一个 add_kernel 在 NVIDIA 和 AMD 上的最终产物差异：

NVIDIA 后端（PTX）
// NVIDIA GPU 上执行的 PTX 汇编 .visible .entry add_kernel( .param .u64 x_ptr, .param .u64 y_ptr, .param .u64 out_ptr, .param .u32 n ) { // 一个 warp (32 线程) 协作 mov.u32 %r1, %ctaid.x; // block ID shl.b32 %r2, %r1, 8; // pid * 256 // 向量化加载：一次 128 字节 ld.global.v4.f32 { %f0,%f1,%f2,%f3 }, [%rd1 + %r2*4]; // 向量化加法 add.f32 %f4, %f0, %f5; add.f32 %f6, %f1, %f7; // 向量化存储 st.global.v4.f32 [%rd3 + %r2*4], { %f4,%f6,%f8,%f10 }; ret; }
AMD 后端（GCN）
// AMD GPU 上执行的 GCN 指令 .amdgpu_kernel add_kernel( .arg x_ptr, 8 .arg y_ptr, 8 .arg out_ptr, 8 .arg n, 4 ) { // 一个 wave (64 线程) 协作 s_load_dwordx4 s[0:3], s[4:5], 0x0 s_load_dword s4, s[4:5], 0x10 s_waitcnt lgkmcnt(0) // 计算偏移 v_mul_i32_i24 v0, v0, 256 // 向量化加载 (buffer_load) buffer_load_dwordx4 v[8:11], v0, s[0:3], 0 // 向量加法 (VOP3) v_add_f32 v8, v8, v12 v_add_f32 v9, v9, v13 // 向量化存储 buffer_store_dwordx4 v[8:11], v0, s[8:11], 0 s_endpgm }
为什么代码不同，但 Python kernel 只有一份
  • TTIR 层：你的 kernel 被翻译成硬件无关的 IR，描述"加载 256 个 float，加法，存储"
  • TTGIR 层：编译器根据目标硬件选择布局编码（NVIDIA 用 warp=32，AMD 用 wave=64）
  • LLVM 层：将 TTGIR 转换成 LLVM IR，然后 LLVM 后端分别生成 PTX（NVIDIA）和 GCN（AMD）
  • 关键：你不需要关心这些差异，Triton 自动处理

源码讲解：布局编码如何适配不同硬件 ——同一个 tl.dot（矩阵乘法），不同硬件用不同的 Tensor Core

# 同一个 Python kernel 中的 tl.dot @triton.jit def matmul_kernel(a_ptr, b_ptr, c_ptr, M, N, K, ...): # ... 加载 a_block, b_block ... c = tl.dot(a_block, b_block) # 矩阵乘法 # --- 编译器在 TTGIR 层的处理 --- # 当目标是 NVIDIA 时，编译器选择： # - MMAEncoding (Tensor Core 布局) # - 16x16x16 矩阵块，适配 H100/A100 的 Tensor Core # - 数据格式：fp16 或 tf32 c = ttg.dot %a, %b, %c { encoding = #ttg.nvidia_mma< version = [2, 0, 0], // Hopper 架构 warpsPerCTA = [4, 2], // 8 个 warp CTALayout = {...} > } : tensor<128 x 128 x f16> -> tensor<128 x 128 x f32> # 当目标是 AMD 时，编译器选择： # - WMMALayout (Matrix Core 布局) # - 32x32x8 矩阵块，适配 MI300X 的 Matrix Core # - 数据格式：fp16 或 bf16 c = ttg.dot %a, %b, %c { encoding = #ttg.amd_wmma< warpsPerCTA = [4, 1], CTALayout = {...} > } : tensor<128 x 128 x f16> -> tensor<128 x 128 x f32> # 当目标是华为昇腾时： # - AscendMmaEncoding (Cube Unit 布局) # - 16x16x16 矩阵块，适配 Da Vinci Cube Unit c = ttg.dot %a, %b, %c { encoding = #ttg.ascend_mma< cubeSize = [16, 16, 16] > } : tensor<128 x 128 x f16> -> tensor<128 x 128 x f32>
布局编码适配的核心逻辑
  • 同一个 tl.dot：Python 层面完全一样
  • TTGIR 层分化：编译器根据目标硬件选择不同的 Encoding（NVIDIA MMA / AMD WMA / Ascend MMA）
  • 硬件匹配：每个 Encoding 对应特定硬件的 Tensor Core / Matrix Core / Cube Unit 规格
  • 性能调优：warpsPerCTA、矩阵块大小等参数需要针对硬件调整

源码讲解：Triton 后端的代码组织 ——每个硬件后端在 Triton 源码中的位置

# Triton 源码目录结构（与硬件适配相关） triton/ ├── lib/ │ ├── Dialect/ │ │ ├── Triton/ # TTIR 定义（硬件无关） │ │ │ ├── IR/ # 操作、类型定义 │ │ │ └── Transform/ # 前端优化 pass │ │ └── TritonGPU/ # TTGIR 定义（GPU 相关） │ │ ├── IR/ # 布局编码：BlockedEncoding, MMAEncoding │ │ └── Transform/ # 后端优化 pass（coalesce, pipeline...） │ └── Conversion/ # TTGIR → LLVM IR 的 lowering │ ├── third_party/ │ ├── nvidia/ # NVIDIA 后端 │ │ ├── lib/ # PTX 生成、warp 级优化 │ │ └── python/ # NVIDIA 特有的 Python API │ │ │ ├── amd/ # AMD 后端 │ │ ├── lib/ # GCN/CDNA 代码生成 │ │ └── python/ # AMD 特有的 API（如 FlashAttention） │ │ │ └── (triton-ascend/) # 华为昇腾后端（独立仓库） │ ├── lib/ # Ascend C 代码生成 │ └── python/ # 昇腾特有 API（tt.fp_to_fp 等）
小结 · 硬件适配的关键
  • TTIR 层：硬件无关，所有后端共享同一套 IR
  • TTGIR 层：通过不同的 Encoding 适配不同硬件的 Tensor Core
  • Backend 层：每个硬件有自己的代码生成器（PTX / GCN / Ascend C）
  • 新增硬件：只需添加新的 Backend 和对应的 Encoding，无需修改用户代码
图 9-B Triton 硬件抽象层
flowchart TD User["用户 Kernel (Python)"] --> TTIR["TTIR (硬件无关)"] TTIR --> TTGIR["TTGIR (布局编码)"] TTGIR --> BN["NVIDIA Backend"] TTGIR --> BA["AMD Backend"] TTGIR --> BH["华为昇腾 Backend"] BN --> PTX["PTX/SASS"] BA --> GCN["GCN/CDNA"] BH --> AscendC["Ascend C"]
9. 版本对比：triton-main vs triton-ascend 版本

目前 Triton 有两个主要版本：triton-main（官方原版，支持 NVIDIA/AMD GPU）和 triton-ascend（华为昇腾适配版，支持 Ascend NPU）。

triton-main（官方原版）

  • 维护者：OpenAI / Triton 社区
  • GitHub：triton-lang/triton
  • 硬件：NVIDIA GPU (CC 8.0+)、AMD GPU (ROCm 6.2+)
  • 安装：pip install triton
  • Python：3.10-3.14
  • 状态：稳定版 v3.x

triton-ascend（昇腾适配版）

  • 维护者：华为昇腾社区
  • GitHub：triton-lang/triton-ascend
  • 硬件：Atlas A2/A3/950 系列 NPU
  • 安装：pip install triton-ascend
  • Python：3.9-3.11
  • 依赖：CANN + TorchNPU
  • 状态：v3.2.2（2026.07.31）
维度triton-maintriton-ascend
编译后端LLVM → PTX/SASS、GCN/CDNALLVM → Ascend C (CANN 工具链)
API 兼容性完整 Triton Python API约 85% API 已支持（持续完善）
内存模型Global → Shared → RegistersGlobal → UB → Registers
特殊 API标准 tl.load/tl.store/tl.dot新增 tt.fp_to_fp (FP8)、scatter_ub_to_out
依赖环境CUDA / ROCmCANN + TorchNPU
DockerDev Container (第三方)官方 Dockerfile (quay.io/ascend/cann)
生态集成PyTorch, vLLM, HuggingFacevLLM-Ascend, sglang-Ascend
迁移指南要点
  • 步骤 1：安装 CANN、TorchNPU，然后 pip install triton-ascend
  • 步骤 2：大部分基础 kernel（使用 tl.load/tl.store/tl.dot）可无缝迁移
  • 步骤 3：涉及特殊内存操作或布局的 kernel 可能需要调整
  • 步骤 4：调整 BLOCK_SIZE 以适配昇腾 Cube Unit 的最优配置
附录 A：源码参考路径 参考

本文档中涉及的代码和概念，对应的 Triton 源码路径如下。读者可参考这些源码深入理解 Triton 的实现细节。

章节内容源码路径
1.1 编译原理TTIR/TTGIR 定义include/triton/Dialect/Triton/IR/
include/triton/Dialect/TritonGPU/IR/
1.1 编译原理布局编码 (BlockedEncoding)include/triton/Dialect/TritonGPU/IR/AttrDefs.td
1.1 编译原理MLIR Dialect 定义include/triton/Dialect/TritonGPU/Transforms/
1.1 编译原理优化 Pass (coalesce, pipeline)lib/Dialect/TritonGPU/Transforms/
2.1 核心 APItl.load / tl.storepython/triton/language/
2.1 核心 APItl.dot (矩阵乘法)python/triton/language/extra/
2.1 核心 APIautotuning 机制python/triton/runtime/autotuner.py
2.1 核心 APItl.constexprpython/triton/language/semantic.py
5. 改写 AttentionFlashAttention kernelthird_party/amd/python/triton/flash_attention.py
third_party/nvidia/lib/
6. 集成autograd.Function 封装python/triton/runtime/
5.1 生态DeepSeek MLA kernelthird_party/ (各后端实现)
9. 硬件适配NVIDIA Backendthird_party/nvidia/
9. 硬件适配AMD Backendthird_party/amd/
9. 硬件适配昇腾 Backendtriton-ascend: lib/ (独立仓库)
10. 版本对比triton-ascend API 扩展triton-ascend: python/triton/
OpenAI Triton · 从介绍到模型实战（完整版） · 穿插编译原理/核心API/生态/硬件适配/版本对比 · 代码块深色高亮 · Mermaid 图独立